/*
 *   код частково запозичено :
 *   http://edu.chpc.ru/cuda/main.html
*/ 
#include<iostream> 
#include <stdio.h> 
#include <math.h> 
#include <stdlib.h> 
 
#define imin(a,b) (a<b?a:b) // обираємо мінімальне з a і b; 
 
const int NSIZE = 2000; // кількість невідомих 
const int threads = 512; // кількість потоків у блоці 
const int blocks = imin(32, ((NSIZE+threads - 1) / threads)); 
// кількість блоків 
 
float *matrA; 
float *vecX, *vecP, *vecR, *vecB, *vecTemp, *vecAP; 
float alpha, beta, normaR; 
float *dev_vec1, *dev_vec2, *dev_matrA; 
// оголошення змінних 
 
void initvector() // ініціалізація векторів 
{ 
    vecTemp = (float *)malloc(NSIZE * sizeof(float)); 
    vecX = (float *) malloc(NSIZE * sizeof(float)); 
    vecP = (float *) malloc(NSIZE * sizeof(float)); 
    vecR = (float *) malloc(NSIZE * sizeof(float)); 
    vecB = (float *) malloc(NSIZE * sizeof(float)); 
    vecAP = (float *) malloc(NSIZE * sizeof(float)); 
    for (int i = 0; i < NSIZE; i++) {// задаємо значення 
        vecX[i] = vecP[i] = vecR[i] = vecTemp[i] = 0.0; //0.0
        vecB[i] = 1.0 / (NSIZE * NSIZE); 
    } 
    // виділення пам’яті для векторів Cuda 
    cudaMalloc( (void**)&dev_vec1, NSIZE*sizeof(float) ) ; 
    cudaMalloc( (void**)&dev_vec2, NSIZE*sizeof(float) ) ; 
} 
 
void initmatrix() { // задання тридіагональної матриці
    matrA = (float *)malloc(NSIZE * NSIZE * sizeof(float)); 
    for (int i = 0; i < NSIZE; i++) 
        for (int j = 0; j < NSIZE; j++) { 
            int k = i * NSIZE + j; 
            matrA[k] = 0; //0 
        } 
    for (int i = 0; i < NSIZE; i++) { 
        matrA[i * NSIZE + i] = 2.0; 
        if (i > 0) { 
            matrA[(i - 1)*NSIZE + i] = -1.0; 
        } 
        if (i < NSIZE - 1) { 
            matrA[(i + 1)*NSIZE + i] = -1.0; 
        } 
    } 
    // виділення пам’яті на відеокарті і копіювання в неї матриці
    cudaMalloc( (void**)&dev_matrA, NSIZE * NSIZE * sizeof(float)); 
    cudaMemcpy(dev_matrA, matrA,  NSIZE * NSIZE * sizeof(float), cudaMemcpyHostToDevice); 
} 
 
// скалярний добуток
float vectorbyvector (float *vect1, float *vect2) 
{ 
    float result; 
    result = 0.0; 
    for (int i = 0; i < NSIZE; i++) { 
        result += vect1[i] * vect2[i]; 
    } 
    return result; 
} 
 
// сума векторів 
float *vectorsum(float *vect1, float *vect2, bool flag) { 
    float *result; 
    result = (float *) malloc(NSIZE * sizeof(float)); 
    for (int i = 0; i < NSIZE; i++) { 
        result[i] = 0.0; 
    } 
    if (flag) { 
        for (int i = 0; i < NSIZE; i++) { 
            result[i] = vect1[i] + vect2[i]; 
        } 
    } 
    else { 
        for (int i = 0; i < NSIZE; i++) { 
            result[i] = vect1[i] - vect2[i]; 
        } 
    } 
    return result; 
} 
 
// множення вектора на число 
float *vectorbynumber(float number, float *vect) { 
    float *result; 
    result = (float *) malloc(NSIZE * sizeof(float)); 
    for (int i = 0; i < NSIZE; i++) { 
        result[i] = 0.0; 
    } 
    for (int i = 0; i < NSIZE; i++) { 
        result[i] += number * vect[i]; 
    } 
    return result; 
} 
__global__ void matdotvec(float *A, float *X, float *Y) { 
    __shared__ float Xds[threads]; // використоаємо розділювану пам’ять 
    int bx = blockIdx.x;    //індекс блоку; 
    int tx = threadIdx.x;   //індекс потоку в блоці; 
    // обчислення індексу 
    int Row = bx * blockDim.x + tx; 
    float Pvalue = 0; 
    for (unsigned int m = 0; m < (NSIZE-1)/threads+1; m++) { 
        if (m*threads + tx< NSIZE) 
        // перевизначення вектору на потрібний діапазон 
             Xds[tx] = X[m*threads + tx]; 
        else 
             Xds[tx] = 0; 
    __syncthreads();// точка синхронізації 
    for (unsigned int k = 0; k < threads; k++) 
        if((Row<NSIZE)&&(m*threads + k < NSIZE)) 
            //збирання множення матриці на вектор всередині блоку
            Pvalue += A[m*threads+Row*NSIZE+k] * Xds[k]; 
    __syncthreads();// точка синхронізації 
    } 
    if(Row < NSIZE) { 
        Y[Row] = Pvalue; // запис у результуючий вектор 
    } 
    __syncthreads();// точка синхронізації 
} 
int main(void) { // головна частина програми
    initvector();// ініціалізація 
    initmatrix();// ініціалізація 
    // копіювання вектора X на відеокарту 
    cudaMemcpy(dev_vec1, vecX, NSIZE * sizeof(float), cudaMemcpyHostToDevice); 
    // Обчислення AX 
    matdotvec<<<blocks,threads>>>(dev_matrA, dev_vec1, dev_vec2); 
    // копіювання з відеокарти на хост
    cudaMemcpy(vecAP, dev_vec2, NSIZE * sizeof(float), cudaMemcpyDeviceToHost); 
    vecR = vecP = vectorsum(vecB, vecAP, false); 
 
    int iter = 0; 
    do { // реалізація ітерацій у методе спряжених градієнтів 
        iter++; 
        // копіювання вектора X на відеокарту 
        cudaMemcpy(dev_vec1, vecP, NSIZE * sizeof(float), cudaMemcpyHostToDevice); 
        // Обчислення AP 
        matdotvec<<<blocks,threads>>>(dev_matrA, dev_vec1, dev_vec2); 
        // копіювання з відеокарти на хост
        cudaMemcpy(vecAP, dev_vec2, NSIZE * sizeof(float), cudaMemcpyDeviceToHost); 
        alpha = vectorbyvector(vecR, vecR) / vectorbyvector (vecP, vecAP); 
        vecX = vectorsum(vecX, vectorbynumber(alpha, vecP), true); 
        vecTemp = vectorsum(vecR, vectorbynumber(alpha, vecAP), false); 
        normaR = sqrt(vectorbyvector(vecR, vecR)); 
        beta = vectorbyvector(vecTemp, vecTemp) / vectorbyvector(vecR, vecR); 
        vecP = vectorsum(vecTemp, vectorbynumber(beta, vecP), true); 
        vecR = vecTemp; 
    } while (normaR> 0.00001);// поки не співпаде по нормі
    FILE *fp; // запис результату в файл 
    fp = fopen("output", "w"); 
    for (int i = 0; i < NSIZE; i++) { 
        fprintf(fp, "%f \n", vecX[i]); 
    } 
    fclose(fp); 
    // звільнення ресурсів cuda
    cudaFree(dev_vec1); 
    cudaFree(dev_vec2);  
    cudaFree(dev_matrA); 
}
